#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pygame, sys
import numpy as np

from scipy.spatial import distance

pygame.init()
W, H = 1200, 700
SCREEN = pygame.display.set_mode((W, H))
pygame.display.set_caption('天地大冲撞')

G = 2
dt = 5

class Body(pygame.sprite.Sprite):
    def __init__(self, position=(80,0), radius=5, color=(255,0,0), mass=1, velocity=0):
        super(Body, self).__init__()
        self.radius= radius
        self.position = position
        self.color = color
        self.mass = mass
        self.velocity = velocity
        self.flag = True

    def draw(self):
        position = int(np.round(self.position[0])), int(np.round(self.position[1]))
        pygame.draw.circle(SCREEN, self.color, position, self.radius)

    def draw_point(self):
        position = int(np.round(self.position[0])), int(np.round(self.position[1]))
        pygame.draw.circle(SCREEN, self.color, position, 1)
 
    def dist(self, other):
        return distance.euclidean(self.position, other.position)

    def direction(self, other):
        if self.dist(other)==0:
            return np.array((0, 0))
        return (np.array(other.position)-np.array(self.position))/self.dist(other)

    def collide(self, other):
        return self.dist(other) < self.radius + other.radius

    def move(self, x):
        position = np.array(self.position, dtype=np.float64)
        position += x
        self.position = position

    def force(self, other):
        if self.dist(other)==0:
            return np.array((0, 0))
        return G * self.mass * other.mass /self.dist(other)**2 * self.direction(other)

    @property
    def total_force(self):
        F = 0.0
        for other in self.groups()[0].sprites():
            if other != self:
                F += self.force(other)
        return F

    def update(self):
        if self.flag:
            F = self.total_force
            a = F / self.mass
            # r = r0 + v dt + a dt^2/2
            self.move(self.velocity * dt + a * dt**2/2)
            self.velocity += a * dt
            if any(self.collide(other) for other in self.groups()[0].sprites() if other !=self):
                self.stop()
            if self.out():
                self.kill()

    def stop(self):
        self.flag = False

    def out(self):
        return self.position[0]<self.radius or self.position[0]>W-self.radius or self.position[1]<self.radius or self.position[1]>H-self.radius


class BodyGroup(pygame.sprite.Group):
    def draw(self):
        for body in self.sprites():
            body.draw()

bodies = (Body(position=(30,80), velocity=np.array((0.,1.)), color=(0,255,0)),
    Body(position=(400,80), velocity=np.array((0.1,0)), radius=30, color=(0,0,255), mass=30),
    Body(position=(300,500), velocity=np.array((-0.1,.3)), color=(0,255,255), radius=6, mass=10),
    Body(position=(600,650), velocity=np.array((.1, -0.1)), color=(100,10,100), radius=8, mass=20),
    Body(position=(80,600), velocity=np.array((.2, -0.2)), color=(10,10,100), radius=13, mass=50))
group = BodyGroup(*bodies)

group.draw()
k=0
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
    if k==5:
        SCREEN.fill((255, 255, 255))
        k=0
    else:
        k += 1
    group.update()
    group.draw()

    pygame.display.update()
    pygame.time.delay(dt)

